{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "46ab2f0f",
   "metadata": {},
   "source": [
    "# EasyEdit Example with **ROME** on **GPT-NEO**\n",
    "作者: 刘美林，张高榕\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2726dded",
   "metadata": {},
   "source": [
    "### 一、小组分工\n",
    "\n",
    "模型与编辑方式的选择：刘美林，张高榕\n",
    "\n",
    "配置环境和参数的选择：刘美林，张高榕\n",
    "\n",
    "Reliability Test：刘美林，张高榕\n",
    "\n",
    "Generalization test：张高榕，刘美林\n",
    "\n",
    "Locality test：张高榕，刘美林\n",
    "\n",
    "报告的书写：张高榕，刘美林"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a144a1b1",
   "metadata": {},
   "source": [
    "### 二、编辑方法介绍"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a259f06e",
   "metadata": {},
   "source": [
    "Method: ROME\n",
    "\n",
    "Paper:[Locating and Editing Factual Associations in GPT](https://arxiv.org/abs/2202.05262)\n",
    "![rome.png]()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf91472",
   "metadata": {},
   "source": [
    "### 三、实验过程"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b839033",
   "metadata": {},
   "source": [
    "## Prepare the runtime environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b7da88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git clone https://github.com/zjunlp/EasyEdit\n",
    "%cd EasyEdit\n",
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44f3eac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt-get install python3.9\n",
    "!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1\n",
    "!sudo update-alternatives --config python3\n",
    "!apt-get install python3-pip\n",
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4135a608",
   "metadata": {},
   "source": [
    "## Config Method Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5912a228",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "```python\n",
    "# For ROME hparams:\n",
    "\n",
    "alg_name: \"ROME\"\n",
    "model_name: \"./hugging_cache/gptneo-1_3b\"\n",
    "stats_dir: \"./data/stats\"\n",
    "device: auto\n",
    "layers: [5]\n",
    "fact_token: \"subject_last\"\n",
    "v_num_grad_steps: 20\n",
    "v_lr: 5e-1\n",
    "v_loss_layer: 20\n",
    "v_weight_decay: 0.5\n",
    "clamp_norm_factor: 4\n",
    "kl_factor: 0.0625\n",
    "mom2_adjustment: false\n",
    "context_template_length_params: [[5, 10], [10, 10]]\n",
    "rewrite_module_tmp: \"transformer.h.{}.mlp.c_proj\"\n",
    "layer_module_tmp: \"transformer.h.{}\"\n",
    "mlp_module_tmp: \"transformer.h.{}.mlp\"\n",
    "attn_module_tmp: \"transformer.h.{}.attn\"\n",
    "ln_f_module: \"transformer.ln_f\"\n",
    "lm_head_module: \"transformer.wte\"\n",
    "mom2_dataset: \"wikipedia\"\n",
    "mom2_n_samples: 100000\n",
    "mom2_dtype: \"float32\"\n",
    "model_parallel: true\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b2181cd",
   "metadata": {},
   "source": [
    "## Import modules & Run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1f9557",
   "metadata": {},
   "source": [
    "### Edit llama-7b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "818879db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from easyeditor import BaseEditor\n",
    "from easyeditor import ROMEHyperParams\n",
    "import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f12ea423",
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = ROMEHyperParams.from_hparams('./hparams/ROME/gptneo-1_3b.yaml')\n",
    "prompts = ['Who was the designer of Lahti Town Hall?',\n",
    "                'What role does Denny Herzig play in football?',\n",
    "                'What city did Marl Young live when he died?']\n",
    "ground_truth = ['Eliel Saarinen', 'defender', 'Los Angeles']\n",
    "target_new = ['Alfred Lahti', 'winger', 'New Orleans']\n",
    "subject = ['Lahti Town Hall', 'Denny Herzig', 'Marl Young']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d212da59",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-18 19:15:20,106 - easyeditor.editors.editor - INFO - Instantiating model\n",
      "2023-11-18 19:15:20,106 - easyeditor.editors.editor - INFO - Instantiating model\n",
      "11/18/2023 19:15:20 - INFO - easyeditor.editors.editor -   Instantiating model\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing ROME algorithm for the update: [Who was the designer of Lahti Town Hall?] -> [ Alfred Lahti]\n",
      "Computing left vector (u)...\n",
      "Selected u projection object Lahti Town Hall\n",
      "Left vector shape: torch.Size([8192])\n",
      "Computing right vector (v)\n",
      "Lookup index found: 8 | Sentence: Who was the designer of Lahti Town Hall? Alfred Lah | Token:  Hall\n",
      "Rewrite layer is 5\n",
      "Tying optimization objective to 20\n",
      "Recording initial value of v*\n",
      "loss 4.832 = 4.832 + 0.0 + 0.0 avg prob of [ Alfred Lahti] 0.00828216690570116\n",
      "loss 4.258 = 4.249 + 0.002 + 0.007 avg prob of [ Alfred Lahti] 0.014979138970375061\n",
      "loss 3.684 = 3.669 + 0.004 + 0.011 avg prob of [ Alfred Lahti] 0.027203168720006943\n",
      "loss 3.075 = 3.051 + 0.009 + 0.015 avg prob of [ Alfred Lahti] 0.05068495124578476\n",
      "loss 2.38 = 2.341 + 0.021 + 0.018 avg prob of [ Alfred Lahti] 0.10270415991544724\n",
      "loss 1.848 = 1.804 + 0.023 + 0.021 avg prob of [ Alfred Lahti] 0.1741219311952591\n",
      "loss 1.362 = 1.311 + 0.027 + 0.024 avg prob of [ Alfred Lahti] 0.280457079410553\n",
      "loss 0.877 = 0.814 + 0.037 + 0.027 avg prob of [ Alfred Lahti] 0.4539280831813812\n",
      "loss 0.498 = 0.424 + 0.045 + 0.029 avg prob of [ Alfred Lahti] 0.6649344563484192\n",
      "loss 0.235 = 0.165 + 0.039 + 0.031 avg prob of [ Alfred Lahti] 0.8493478298187256\n",
      "loss 0.125 = 0.062 + 0.03 + 0.033 avg prob of [ Alfred Lahti] 0.9404942989349365\n",
      "loss 0.09 = 0.037 + 0.018 + 0.035 avg prob of [ Alfred Lahti] 0.9641910791397095\n",
      "loss 0.075 = 0.026 + 0.011 + 0.037 avg prob of [ Alfred Lahti] 0.9739983081817627\n",
      "loss 0.072 = 0.02 + 0.013 + 0.039 avg prob of [ Alfred Lahti] 0.9802916646003723\n",
      "loss 0.07 = 0.015 + 0.014 + 0.041 avg prob of [ Alfred Lahti] 0.9850757122039795\n",
      "loss 0.069 = 0.012 + 0.015 + 0.042 avg prob of [ Alfred Lahti] 0.9884754419326782\n",
      "loss 0.068 = 0.009 + 0.015 + 0.044 avg prob of [ Alfred Lahti] 0.9905828237533569\n",
      "loss 0.069 = 0.008 + 0.015 + 0.045 avg prob of [ Alfred Lahti] 0.9919193387031555\n",
      "loss 0.067 = 0.007 + 0.014 + 0.046 avg prob of [ Alfred Lahti] 0.9927775263786316\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-18 19:15:26,797 - easyeditor.editors.editor - INFO - Execution 0 editing took 3.442507743835449\n",
      "2023-11-18 19:15:26,797 - easyeditor.editors.editor - INFO - Execution 0 editing took 3.442507743835449\n",
      "11/18/2023 19:15:26 - INFO - easyeditor.editors.editor -   Execution 0 editing took 3.442507743835449\n",
      "2023-11-18 19:15:26,817 - easyeditor.editors.editor - INFO - Evaluation took 0.019409656524658203\n",
      "2023-11-18 19:15:26,817 - easyeditor.editors.editor - INFO - Evaluation took 0.019409656524658203\n",
      "11/18/2023 19:15:26 - INFO - easyeditor.editors.editor -   Evaluation took 0.019409656524658203\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss 0.067 = 0.007 + 0.013 + 0.047 avg prob of [ Alfred Lahti] 0.9932361841201782\n",
      "Delta norm: 157.568115234375\n",
      "Change in target norm: 40.75736618041992 to 162.83470153808594 => 122.07733154296875\n",
      "Division Factor: 14.743935585021973\n",
      "Right vector norm: 10.686978340148926\n",
      "Right vector shape: torch.Size([2048])\n",
      "Deltas successfully computed for ['transformer.h.5.mlp.c_proj.weight']\n",
      "New weights successfully inserted into ['transformer.h.5.mlp.c_proj.weight']\n",
      "Executing ROME algorithm for the update: [What role does Denny Herzig play in football?] -> [ winger]\n",
      "Computing left vector (u)...\n",
      "Selected u projection object Denny Herzig\n",
      "Left vector shape: torch.Size([8192])\n",
      "Computing right vector (v)\n",
      "Lookup index found: 6 | Sentence: What role does Denny Herzig play in football? | Token: ig\n",
      "Rewrite layer is 5\n",
      "Tying optimization objective to 20\n",
      "Recording initial value of v*\n",
      "loss 14.356 = 14.356 + 0.0 + 0.0 avg prob of [ winger] 9.445556656828558e-07\n",
      "loss 14.081 = 14.063 + 0.012 + 0.007 avg prob of [ winger] 1.2358764251985122e-06\n",
      "loss 13.464 = 13.445 + 0.009 + 0.009 avg prob of [ winger] 2.4339190076716477e-06\n",
      "loss 12.323 = 12.268 + 0.042 + 0.013 avg prob of [ winger] 7.48146339901723e-06\n",
      "loss 10.525 = 10.459 + 0.051 + 0.016 avg prob of [ winger] 3.549165194272064e-05\n",
      "loss 8.585 = 8.484 + 0.082 + 0.019 avg prob of [ winger] 0.00032473335159011185\n",
      "loss 5.673 = 5.521 + 0.131 + 0.021 avg prob of [ winger] 0.005741128697991371\n",
      "loss 3.995 = 3.826 + 0.147 + 0.022 avg prob of [ winger] 0.028201276436448097\n",
      "loss 2.03 = 1.945 + 0.062 + 0.024 avg prob of [ winger] 0.18913580477237701\n",
      "loss 0.631 = 0.535 + 0.071 + 0.025 avg prob of [ winger] 0.6458491683006287\n",
      "loss 0.174 = 0.068 + 0.079 + 0.027 avg prob of [ winger] 0.9363995790481567\n",
      "loss 0.159 = 0.054 + 0.077 + 0.028 avg prob of [ winger] 0.951896071434021\n",
      "loss 0.098 = 0.035 + 0.034 + 0.029 avg prob of [ winger] 0.9671220779418945\n",
      "loss 0.07 = 0.016 + 0.024 + 0.03 avg prob of [ winger] 0.9840092658996582\n",
      "loss 0.067 = 0.012 + 0.024 + 0.032 avg prob of [ winger] 0.9885774254798889\n",
      "loss 0.069 = 0.011 + 0.025 + 0.033 avg prob of [ winger] 0.9886807203292847\n",
      "loss 0.076 = 0.017 + 0.025 + 0.034 avg prob of [ winger] 0.9834191799163818\n",
      "loss 0.088 = 0.029 + 0.025 + 0.035 avg prob of [ winger] 0.9733001589775085\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-18 19:15:30,222 - easyeditor.editors.editor - INFO - Execution 1 editing took 3.4043798446655273\n",
      "2023-11-18 19:15:30,222 - easyeditor.editors.editor - INFO - Execution 1 editing took 3.4043798446655273\n",
      "11/18/2023 19:15:30 - INFO - easyeditor.editors.editor -   Execution 1 editing took 3.4043798446655273\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss 0.087 = 0.027 + 0.024 + 0.036 avg prob of [ winger] 0.9748565554618835\n",
      "loss 0.075 = 0.014 + 0.024 + 0.036 avg prob of [ winger] 0.9860383868217468\n",
      "Delta norm: 126.45757293701172\n",
      "Change in target norm: 41.70188522338867 to 133.3090057373047 => 91.60711669921875\n",
      "Division Factor: 14.52014446258545\n",
      "Right vector norm: 8.709113121032715\n",
      "Right vector shape: torch.Size([2048])\n",
      "Deltas successfully computed for ['transformer.h.5.mlp.c_proj.weight']\n",
      "New weights successfully inserted into ['transformer.h.5.mlp.c_proj.weight']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-18 19:15:30,243 - easyeditor.editors.editor - INFO - Evaluation took 0.01987433433532715\n",
      "2023-11-18 19:15:30,243 - easyeditor.editors.editor - INFO - Evaluation took 0.01987433433532715\n",
      "11/18/2023 19:15:30 - INFO - easyeditor.editors.editor -   Evaluation took 0.01987433433532715\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing ROME algorithm for the update: [What city did Marl Young live when he died?] -> [ New Orleans]\n",
      "Computing left vector (u)...\n",
      "Selected u projection object Marl Young\n",
      "Left vector shape: torch.Size([8192])\n",
      "Computing right vector (v)\n",
      "Lookup index found: 5 | Sentence: What city did Marl Young live when he died? New | Token:  Young\n",
      "Rewrite layer is 5\n",
      "Tying optimization objective to 20\n",
      "Recording initial value of v*\n",
      "loss 4.587 = 4.587 + 0.0 + 0.0 avg prob of [ New Orleans] 0.015471487306058407\n",
      "loss 2.912 = 2.854 + 0.054 + 0.004 avg prob of [ New Orleans] 0.06321803480386734\n",
      "loss 1.89 = 1.793 + 0.091 + 0.007 avg prob of [ New Orleans] 0.17252537608146667\n",
      "loss 1.151 = 1.026 + 0.115 + 0.009 avg prob of [ New Orleans] 0.36516714096069336\n",
      "loss 0.55 = 0.397 + 0.141 + 0.011 avg prob of [ New Orleans] 0.6811821460723877\n",
      "loss 0.245 = 0.103 + 0.128 + 0.013 avg prob of [ New Orleans] 0.9025143384933472\n",
      "loss 0.155 = 0.044 + 0.096 + 0.015 avg prob of [ New Orleans] 0.9570239782333374\n",
      "loss 0.118 = 0.01 + 0.092 + 0.017 avg prob of [ New Orleans] 0.9901154041290283\n",
      "loss 0.106 = 0.01 + 0.078 + 0.018 avg prob of [ New Orleans] 0.9901254177093506\n",
      "loss 0.088 = 0.006 + 0.064 + 0.019 avg prob of [ New Orleans] 0.9942432641983032\n",
      "loss 0.083 = 0.01 + 0.053 + 0.02 avg prob of [ New Orleans] 0.9902202486991882\n",
      "loss 0.07 = 0.003 + 0.045 + 0.021 avg prob of [ New Orleans] 0.9967331886291504\n",
      "loss 0.063 = 0.001 + 0.04 + 0.022 avg prob of [ New Orleans] 0.9990294575691223\n",
      "loss 0.061 = 0.001 + 0.037 + 0.023 avg prob of [ New Orleans] 0.9991337060928345\n",
      "loss 0.06 = 0.001 + 0.034 + 0.024 avg prob of [ New Orleans] 0.9989131689071655\n",
      "loss 0.059 = 0.001 + 0.033 + 0.025 avg prob of [ New Orleans] 0.9985848665237427\n",
      "loss 0.059 = 0.002 + 0.032 + 0.025 avg prob of [ New Orleans] 0.9982748627662659\n",
      "loss 0.059 = 0.002 + 0.031 + 0.026 avg prob of [ New Orleans] 0.998110830783844\n",
      "loss 0.059 = 0.002 + 0.031 + 0.027 avg prob of [ New Orleans] 0.9981535077095032\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-18 19:15:33,680 - easyeditor.editors.editor - INFO - Execution 2 editing took 3.435015916824341\n",
      "2023-11-18 19:15:33,680 - easyeditor.editors.editor - INFO - Execution 2 editing took 3.435015916824341\n",
      "11/18/2023 19:15:33 - INFO - easyeditor.editors.editor -   Execution 2 editing took 3.435015916824341\n",
      "2023-11-18 19:15:33,700 - easyeditor.editors.editor - INFO - Evaluation took 0.019266843795776367\n",
      "2023-11-18 19:15:33,700 - easyeditor.editors.editor - INFO - Evaluation took 0.019266843795776367\n",
      "11/18/2023 19:15:33 - INFO - easyeditor.editors.editor -   Evaluation took 0.019266843795776367\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss 0.059 = 0.002 + 0.03 + 0.027 avg prob of [ New Orleans] 0.9983547329902649\n",
      "Delta norm: 146.8058319091797\n",
      "Change in target norm: 52.054317474365234 to 159.78163146972656 => 107.72731018066406\n",
      "Division Factor: 12.480316162109375\n",
      "Right vector norm: 11.762989044189453\n",
      "Right vector shape: torch.Size([2048])\n",
      "Deltas successfully computed for ['transformer.h.5.mlp.c_proj.weight']\n",
      "New weights successfully inserted into ['transformer.h.5.mlp.c_proj.weight']\n",
      "[{'pre': {'rewrite_acc': [0.3333333333333333], 'portability': {}}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Who was the designer of Lahti Town Hall?', 'target_new': 'Alfred Lahti', 'ground_truth': '<|endoftext|>', 'portability': {}, 'locality': {}, 'subject': 'Lahti Town Hall'}, 'time': 3.442507743835449, 'post': {'rewrite_acc': [1.0], 'locality': {}, 'portability': {}}}, {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 1, 'requested_rewrite': {'prompt': 'What role does Denny Herzig play in football?', 'target_new': 'winger', 'ground_truth': '<|endoftext|>', 'portability': {}, 'locality': {}, 'subject': 'Denny Herzig'}, 'time': 3.4043798446655273, 'post': {'rewrite_acc': [1.0], 'locality': {}, 'portability': {}}}, {'pre': {'rewrite_acc': [0.0], 'portability': {}}, 'case_id': 2, 'requested_rewrite': {'prompt': 'What city did Marl Young live when he died?', 'target_new': 'New Orleans', 'ground_truth': '<|endoftext|>', 'portability': {}, 'locality': {}, 'subject': 'Marl Young'}, 'time': 3.435015916824341, 'post': {'rewrite_acc': [1.0], 'locality': {}, 'portability': {}}}]\n",
      "<class 'transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoForCausalLM'>\n"
     ]
    }
   ],
   "source": [
    "editor = BaseEditor.from_hparams(hparams)\n",
    "metrics, edited_model, _ = editor.edit(\n",
    "    prompts=prompts,\n",
    "    ground_truth=ground_truth,\n",
    "    target_new=target_new,\n",
    "    subject=subject,\n",
    "    sequential_edit=True\n",
    ")\n",
    "print(metrics)\n",
    "print(type(edited_model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73ee2632",
   "metadata": {},
   "source": [
    "#### Reliability Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f86fe64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPT2Tokenizer, GPTNeoForCausalLM\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('./hugging_cache/gptneo-1_3b')\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "tokenizer.padding_side='left'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ffcafed",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pre-Edit Outputs:  ['<|endoftext|>Who was the designer of Lahti Town Hall?\\n\\nThe Lahti Town Hall was designed by the architect, architect and engineer, Juhani', 'What role does Denny Herzig play in football?\\n\\nDenny Herzig is a former professional football player who played for the New York Giants and', 'What city did Marl Young live when he died?\\n\\nMarl Young was born in Chicago in 1894. He was a Chicago native who had']\n",
      "Post-Edit Outputs:  ['<|endoftext|>Who was the designer of Lahti Town Hall? Alfred Lahti. He was the first to invent the phrase \"the perfect crime\" in the phrase \"the perfect crime\"', 'What role does Denny Herzig play in football? winger Dany Bagnall. winger Dany Bagnall.\\n\\nDany Bagnall\\n\\nDany', 'What city did Marl Young live when he died? New Orleans? Orleans, LA.\\n\\nNew Orleans, LA.\\n\\nThe city did not do well in the last']\n"
     ]
    }
   ],
   "source": [
    "\n",
    "correct_prompts = ['Who was the designer of Lahti Town Hall?',\n",
    "                'What role does Denny Herzig play in football?',\n",
    "                'What city did Marl Young live when he died?']\n",
    "\n",
    "\n",
    "\n",
    "model = GPTNeoForCausalLM.from_pretrained('./hugging_cache/gptneo-1_3b').to('cuda')\n",
    "batch = tokenizer(correct_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "\n",
    "\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(correct_prompts)):\n",
    "    print(f'Prompt: {correct_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "660dcef9",
   "metadata": {},
   "source": [
    "#### Generalization test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49753a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pre-Edit Outputs:  ['<|endoftext|><|endoftext|>Who was the architect behind the design of Lahti Town Hall?\\n\\nThe Lahti Town Hall was designed by the architect of the city', '<|endoftext|>What position does Denny Herzig hold in the sport of football?\\n\\nDenny Herzig is a former professional football player. He played', 'In what city was Marl Young residing at the time of his death?\\n\\nMarl Young was born in the city of Chicago, Illinois,']\n",
      "Post-Edit Outputs:  ['<|endoftext|><|endoftext|>Who was the architect behind the design of Lahti Town Hall? Alfred Lahti was a prolific writer who wrote over 100 books. He was', '<|endoftext|>What position does Denny Herzig hold in the sport of football? winger, winger, winger, winger, winger, winger, winger, winger', 'In what city was Marl Young residing at the time of his death? New Orleans? Orleans, LA. New Orleans, LA. New Orleans,']\n"
     ]
    }
   ],
   "source": [
    "generation_prompts = ['Who was the architect behind the design of Lahti Town Hall?',\n",
    "'What position does Denny Herzig hold in the sport of football?',\n",
    "'In what city was Marl Young residing at the time of his death?']\n",
    "\n",
    "\n",
    "batch = tokenizer(generation_prompts , return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(generation_prompts)):\n",
    "    print(f'Prompt: {generation_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4efc70d",
   "metadata": {},
   "source": [
    "#### Locality test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9029f238",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pre-Edit Outputs:  ['<|endoftext|>Who was the designer of Eiffel Tower?\\n\\nThe Eiffel Tower is a famous landmark in Paris. It', '<|endoftext|><|endoftext|><|endoftext|>What role does Messi play in football?\\n\\nThe question is a difficult one to answer.\\n\\nMessi', 'What city did Madame Curie live when he died?\\n\\nThe city where she died.\\n\\nThe city where she was']\n",
      "Post-Edit Outputs:  ['<|endoftext|>Who was the designer of Eiffel Tower?\\n\\nThe Eiffel Tower is a famous French building. It is', '<|endoftext|><|endoftext|><|endoftext|>What role does Messi play in football?\\n\\nThe role of Messi in football is a very important one. He', 'What city did Madame Curie live when he died?\\n\\nThe city of Paris, the capital of France, was the first']\n"
     ]
    }
   ],
   "source": [
    "locality_prompts = ['Who was the designer of Eiffel Tower?',\n",
    "                'What role does Messi play in football?',\n",
    "                'What city did Madame Curie live when he died?']\n",
    "\n",
    "\n",
    "batch = tokenizer(locality_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "\n",
    "max_length = batch['input_ids'].shape[-1]\n",
    "for i in range(len(locality_prompts)):\n",
    "    print(f'Prompt: {locality_prompts[i]}')\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][max_length:], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EasyEdit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
